import math
import functools
import numpy as np
import quimb as qu
from quimb.core import make_immutable

@functools.lru_cache(16)
def basic_ptm_vec(i, dim, **kwargs):
	shape = (dim, 1)
	ptm_x = np.zeros(shape, dtype=complex)
	ptm_x[i] = 1.0
	return qu.qu(ptm_x, **kwargs)

@functools.lru_cache(8)
def ptm_xplus(**kwargs):
	r"""Returns PTM representation of plus-state, aka. ``|+><+|``, +X eigenstate"""
	return qu.qu([[1.0], [1.0], [0.0], [0.0]], **kwargs) / math.sqrt(2)

@functools.lru_cache(8)
def ptm_xminus(**kwargs):
	r"""Returns PTM representation of minus-state, aka. ``|-><-|``, -X eigenstate"""
	return qu.qu([[1.0], [-1.0], [0.0], [0.0]], **kwargs) / math.sqrt(2)

@functools.lru_cache(8)
def ptm_yplus(**kwargs):
	r"""Returns PTM representation of yplus-state, aka. ``|y+><y+|``, +Y eigenstate"""
	return qu.qu([[1.0], [0.0], [1.0], [0.0]], **kwargs) / math.sqrt(2)

@functools.lru_cache(8)
def ptm_yminus(**kwargs):
	r"""Returns PTM representation of yminus-state, aka. ``|y-><y-|``, -Y eigenstate"""
	return qu.qu([[1.0], [0.0], [-1.0], [0.0]], **kwargs) / math.sqrt(2)

@functools.lru_cache(8)
def ptm_zplus(**kwargs):
	r"""Returns PTM representation of up-state, aka. ``|0><0|``, +Z eigenstate"""
	return qu.qu([[1.0], [0.0], [0.0], [1.0]], **kwargs) / math.sqrt(2)

@functools.lru_cache(8)
def ptm_zminus(**kwargs):
	r"""Returns PTM representation of down-state, aka. ``|1><1|``, -Z eigenstate"""
	return qu.qu([[1.0], [0.0], [0.0], [-1.0]], **kwargs) / math.sqrt(2)

ptm_plus = ptm_xplus
ptm_minus = ptm_xminus

ptm_up = ptm_zplus
ptm_down = ptm_zminus

@functools.lru_cache(16)
def ptm_pauli_proj(xyz_proj, **kwargs):
	xyz_proj_map = {0: 'x+', 'x+': 'x+', 'X+': 'x+',
	                1: 'x-', 'x-': 'x-', 'X-': 'x-',
	                2: 'y+', 'y+': 'y+', 'Y+': 'y+',
	                3: 'y-', 'y-': 'y-', 'Y-': 'y-',
	                4: 'z+', 'z+': 'z+', 'Z+': 'z+',
	                5: 'z-', 'z-': 'z-', 'Z-': 'z-', }
	ptm_op_map = {
		'x+': ptm_xplus,
		'x-': ptm_xminus,
		'y+': ptm_yplus,
		'y-': ptm_yminus,
		'z+': ptm_zplus,
		'z-': ptm_zminus,
	}
	proj = ptm_op_map[xyz_proj_map.get(xyz_proj, 'z+')](**kwargs)
	make_immutable(proj)
	return proj

def ptm_bloch_state(ax, ay, az, purified=False, **kwargs):
	n = (ax ** 2 + ay ** 2 + az ** 2) ** 0.5
	if purified:
		ax, ay, az = (a / n for a in (ax, ay, az))
	return qu.qu([[1.0], [ax], [ay], [az]], **kwargs) / math.sqrt(2)
